/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANSFER_SHAPE_ACCORDING_TO_FORMAT_H_
#define FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANSFER_SHAPE_ACCORDING_TO_FORMAT_H_

#include <functional>
#include <vector>
#include "common/fe_inner_error_codes.h"
#include "common/fe_utils.h"
#include "common/format/axis_util.h"
#include "common/util/op_info_util.h"
#include "graph/compute_graph.h"

namespace fe {

/* The first parameter is axis value, second is new shape and third is
 * op implementation type. */
using GetNewShapeByAxisValueAndFormat =
    std::function<Status(ge::GeShape &, const int64_t &, vector<int64_t> &, vector<int64_t> &)>;

using GetNewShapeByAxisValueAndFormatPtr = std::shared_ptr<GetNewShapeByAxisValueAndFormat>;

struct ShapeAndFormatInfo {
  const ge::GeShape &old_shape;
  ge::GeShape &new_shape;
  const ge::Format &old_format;
  const ge::Format &new_format;
  const ge::DataType &current_data_type;
  const int64_t &op_impl_type;
  const int64_t &group_count;
};

using ShapeAndFormat = struct ShapeAndFormatInfo;

uint32_t GetC0(ge::DataType dtype);

class ShapeTransferAccordingToFormat {
 public:
  ShapeTransferAccordingToFormat();

  ~ShapeTransferAccordingToFormat();

  ShapeTransferAccordingToFormat(const ShapeTransferAccordingToFormat &) = delete;

  ShapeTransferAccordingToFormat &operator=(const ShapeTransferAccordingToFormat &) = delete;

  static Status GetShapeAccordingToFormat(ShapeAndFormat &input_and_output_info, int64_t *c = nullptr);

  /*----------Below is the function of getting new shape----------------------*/
  static Status GetNCHWShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                        const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetNHWCShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                        const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetNC1HWC0ShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                           const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetFzShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                      const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetFzC04ShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                         const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetHWCNShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                        const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetCHWNShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                        const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetC1HWNCoC0ShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                             const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetNzShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                      const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetNDC1HWC0ShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                            const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetFz3DShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                        const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetFz3DTransposeShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                                 const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

  static Status GetFzLstmShapeByAxisValue(ge::GeShape &new_shape, const int64_t &impl_type,
                                          const vector<int64_t> &axis_value, const vector<int64_t> &nd_value);

 private:
  /* map of GetAxisValueInfoByFormat, get axis value by different original
   * formats. */
  static const std::map<ge::Format, GetNewShapeByAxisValueAndFormatPtr> get_new_shape_func_map;

  static int64_t GetAsisEnlargeValue(const int64_t& cin, const int64_t& cout, const int64_t& c0, const int64_t& group);
};
}  // namespace fe

#endif  // FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_SHAPE_FORMAT_TRANSFER_TRANSFER_SHAPE_ACCORDING_TO_FORMAT_H_
